import torch
import cv2


def load_model(model_path):
    """
    加载预训练的YOLOv9模型。
    :param model_path: 模型文件的路径
    :return: 加载的模型
    """
    model = torch.hub.load('ultralytics/yolov9', 'custom', path=model_path, force_reload=True)
    model.eval()
    return model


def draw_boxes(results, frame):
    """
    在视频帧上绘制检测到的边界框和标签。
    :param results: 模型的预测结果
    :param frame: 视频帧
    :return: 标注后的视频帧
    """
    labels = results.names
    for result in results.xyxy[0]:  # 对每个检测对象进行操作
        x1, y1, x2, y2, conf, cls_id = map(int, result[:6])
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 画边界框
        cv2.putText(frame, f"{labels[cls_id]} {conf:.2f}",
                    (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
                    0.5, (0, 255, 0), 2)
    return frame
